import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

#在预训练模型上进行微调
# load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")
# replace the classifier with a new one, that has
# num_classes which is user-defined
num_classes = 2  # 1 class (person) + background
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

def get_model_instance_segmentation(num_classes):
    #加载一个预训练的实例分割模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

    #获取分类器里面的输入特征数量
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    #替换一个新的预训练头
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features,num_classes)

    #获取掩码分类器里面的输入特征数量
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidder_layer = 256

    #替换掩码预测器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask,
        hidder_layer,
        num_classes
    )

    return model

